package com.lee.algorithm.array;

import java.util.Comparator;

/**
 * @Description: 堆
 *      不传比较器，默认升序排序，即小顶堆
 * @Author: 青石路
 * @CreateDate: 2022/4/10 9:59
 **/
public class Heap<E> {

    private static final int DEFAULT_INITIAL_CAPACITY = 11;

    private Object[] arr;

    private final Comparator<? super E> comparator;

    private int size = 0;

    public Heap() {
        this(DEFAULT_INITIAL_CAPACITY, null);
    }

    public Heap(int initialCapacity) {
        this(initialCapacity, null);
    }

    public Heap(int initialCapacity, Comparator<? super E> comparator) {
        if (initialCapacity < 1)
            throw new IllegalArgumentException();
        this.arr = new Object[initialCapacity];
        this.comparator = comparator;
    }

    /**
     * 在位置 k 处插入元素 x，逐层将 x 往根上提升，保持堆属性不变
     * @param k
     * @param x
     * @author 博客园 @ 青石路
     */
    private void shiftUp(int k, E x) {
        if (comparator != null) {
            // 按自定义比较器排序处理
            shiftUpUsingComparator(k, x);
        } else {
            // 按自然排序处理
            shiftUpComparable(k, x);
        }
    }

    /**
     * 在位置 k 处插入元素 x，然后逐层将 x 往根上提升，直至最终满足堆属性
     *      自定义比较器进行比较
     * @param k
     * @param x
     * @author 博客园 @ 青石路
     */
    private void shiftUpUsingComparator(int k, E x) {
        while (k > 0) {
            // >>>：无符号右移
            // 找到 k 位置的父位置：parent
            int parent = (k - 1) >>> 1;
            // 元素 x 的父元素：e
            Object e = arr[parent];
            if (comparator.compare(x, (E) e) >= 0) {
                // x 与 e 的大小关系符合自定义比较器的规则，说明符合堆属性
                break;
            }

            // x 与 e 的大小关系不符合自定义比较器的规则，需要进行调整
            arr[k] = e;
            k = parent;
        }
        arr[k] = x;
    }

    /**
     * 在位置 k 处插入元素 x，然后逐层将 x 往根上移动（上移），直至最终满足堆属性
     *       利用 x 的自然比较器进行比较，默认升序
     *
     * @param k
     * @param x
     * @author 博客园 @ 青石路
     */
    private void shiftUpComparable(int k, E x) {
        Comparable<? super E> key = (Comparable<? super E>) x;
        while (k > 0) {
            // >>>：无符号右移
            // 找到 k 位置的父位置：parent
            int parent = (k - 1) >>> 1;
            // 元素 x 的父元素：e
            Object e = arr[parent];
            if (key.compareTo((E) e) >= 0) {
                // x 与 e 的大小关系符合自然比较器的规则，说明符合堆属性
                break;
            }

            // x 与 e 的大小关系不符合自然比较器的规则，需要进行调整
            arr[k] = e;
            k = parent;
        }
        arr[k] = key;
    }

    /**
     * 在位置 k 处插入元素 x，然后逐层将 x 往叶子上移动（下移），直至满足堆属性
     * @param k
     * @param x
     * @author 博客园 @ 青石路
     */
    private void shiftDown(int k, E x) {
        if (comparator != null) {
            // 按自定义比较器排序处理
            shiftDownUsingComparator(k, x);
        } else {
            // 按自然排序处理
            shiftDownComparable(k, x);
        }
    }

    private void shiftDownComparable(int k, E x) {
        Comparable<? super E> key = (Comparable<? super E>)x;
        int half = size >>> 1;        // 非叶子节点索引的最大值
        while (k < half) {
            int child = (k << 1) + 1; // 左孩子索引
            Object c = arr[child];
            int right = child + 1;    // 右孩子索引
            if (right < size &&
                    ((Comparable<? super E>) c).compareTo((E) arr[right]) > 0) {
                c = arr[child = right];
            }
            if (key.compareTo((E) c) <= 0) {
                // 此时满足堆属性，x 无需再移动
                break;
            }
            arr[k] = c;
            k = child;
        }
        arr[k] = key;
    }

    private void shiftDownUsingComparator(int k, E x) {
        int half = size >>> 1;          // 非叶子节点索引的最大值
        while (k < half) {
            int child = (k << 1) + 1;   // 左孩子索引
            Object c = arr[child];
            int right = child + 1;      // 右孩子索引
            if (right < size &&
                    comparator.compare((E) c, (E) arr[right]) > 0) {
                // c 记录的是左右孩子中的大者（大顶堆）或小者（小顶堆）
                c = arr[child = right];
            }
            if (comparator.compare(x, (E) c) <= 0) {
                // 此时满足堆属性，x 无需再移动
                break;
            }

            // 不满足堆属性，k 继续往下移动
            arr[k] = c;
            k = child;
        }
        arr[k] = x;
    }

    /**
     * 插入
     * @param e
     * @author 博客园 @ 青石路
     */
    public void insert(E e) {
        if (e == null)
            throw new NullPointerException();
        int i = size;
        if (i >= arr.length) {
            // TODO 扩容处理：arr空间扩大，旧元素拷贝到新数组中
        }
        size = i + 1;
        if (i == 0) {
            arr[0] = e;
        } else {
            shiftUp(i, e);
        }
    }

    /**
     * 查找元素的位置索引
     *      时间复杂度：O(N)
     * @param o
     * @return
     * @author 博客园 @ 青石路
     */
    public int indexOf(Object o) {
        if (o != null) {
            for (int i = 0; i < size; i++) {
                if (o.equals(arr[i])) {
                    return i;
                }
            }
        }
        return -1;
    }

    /**
     * 移除根节点元素
     * @return
     * @author 博客园 @ 青石路
     */
    public boolean remove(E o) {
        int i = indexOf(o);
        if (i == -1) {
            return false;
        } else {
            removeAt(i);
            return true;
        }
    }

    /**
     * 删除指定位置的元素
     * @param i
     * @return
     * @author 博客园 @ 青石路
     */
    public E removeAt(int i) {
        int s = --size;
        if (s == i) {
            arr[i] = null;
        } else {
            // 获取最后一个元素
            E moved = (E) arr[s];
            arr[s] = null;
            shiftDown(i, moved);
            if (arr[i] == moved) {
                shiftUp(i, moved);
                if (arr[i] != moved) {
                    return moved;
                }
            }
        }
        return null;
    }

    /**
     * 获取（大顶堆的）最大值或（小顶堆的）最小值
     * @return
     * @author 博客园 @ 青石路
     */
    public E peek() {
        return (size == 0) ? null : (E) arr[0];
    }

    /**
     * 元素替换
     * @param i 位置索引
     * @param o 目标元素
     * @return 替换前的旧元素
     * @author 博客园 @ 青石路
     */
    public E replace(int i, E o) {
        if (i <= 0 || i >= size) {
            return null;
        }
        E old = (E) arr[i];
        shiftDown(i, o);
        if (arr[i] == o) {
            shiftUp(i, o);
            if (arr[i] != o) {
                return old;
            }
        }
        return null;
    }

    /**
     * 构建堆
     * @param initArr
     * @author 博客园 @ 青石路
     */
    public void buildHeap(E[] initArr) {
        if (initArr == null) {
            return;
        }
        for (int i=0; i<initArr.length; i++) {
            insert(initArr[i]);
        }
    }
}
